import torch
import re
from dataset.vocab import Vocab
from torch import nn
import time
from backbones.transformer import ChatNet

if __name__ == '__main__':
    sentence = "Where is the captain of china?\t<bos> The captain of china is beijing.<eos>"
    s1, s2 = sentence.split("\t")
    # 词元化(是否是同一个词表)
    input_tokens = [re.sub(r"[^a-zA-Z<>]+", " ", s1).lower().strip().split()]
    output_tokens = [re.sub(r"[^a-zA-Z<>]+", " ", s2).lower().strip().split()]
    tokens = input_tokens + output_tokens

    vocab = Vocab(tokens, 0)  # 英文统一使用同一个词元表

    # 将输入信息进行idx转化
    encode_idx = [vocab.to_idx(line) for line in input_tokens]
    decode_idx = [vocab.to_idx(line) for line in output_tokens]
    encode_idx = torch.tensor(encode_idx)  # (1,6)
    decode_idx = torch.tensor(decode_idx)

    # 解码器中的序列设置 （此处为teach force 教师强制模式）
    decode_input = decode_idx[:, :-1]  # (1,7)
    decode_output = decode_idx[:, 1:]  # (1,7)

    vocab_size = len(vocab)
    device = torch.device("cuda")

    model = ChatNet(vocab_size, 10, 2, 1, 1, device=device)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    epochs = 1000
    for epoch in range(epochs):
        optimizer.zero_grad()
        inputs1 = encode_idx.to(device)
        inputs2 = decode_input.to(device)

        predicts = model(inputs1, inputs2)

        loss = criterion(predicts.reshape(-1, vocab_size), decode_output.to(device).reshape(-1))
        loss.backward()

        # 梯度裁剪
        nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
        optimizer.step()

        if (epoch + 1) % 100 == 0:
            print(f"epoch:{epoch + 1}/{epochs} -- loss:{loss.item():.4f}")

    # ===================== 预测 =====================
    model.eval()
    outputs = torch.tensor([[vocab.to_idx("<bos>")]]).to(device)  # (1,1)
    while True:
        results = model(encode_idx.to(device), outputs)
        idx = torch.argmax(results[:, -1:], dim=-1)
        word = vocab[idx.squeeze(0).item()]
        time.sleep(0.5)
        print(word, end=" ")
        if word == "<eos>":
            break
        outputs = torch.cat([outputs, idx], dim=-1)
